import sys
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                               QHBoxLayout, QLabel, QPushButton, QRadioButton,
                               QButtonGroup, QCheckBox, QLineEdit, QProgressBar,
                               QFileDialog, QMessageBox)
from PySide6.QtCore import Qt, QThread, Signal
import os

from gui.model_loader import ModelLoader
from gui.ocr_worker import OCRWorker


class StartupWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model_loader = ModelLoader()
        self.model_loader.loading_finished.connect(self.on_model_loading_finished)
        self.model_loader.start()  # 启动后台加载模型

        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("身份证OCR工具")
        self.setMinimumSize(550, 450)

        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(30, 30, 30, 30)

        # 第一行：软件名称
        title_label = QLabel("身份证OCR工具")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setStyleSheet("font-size: 24px; font-weight: bold; margin: 20px 0;")
        main_layout.addWidget(title_label)

        # 添加模型加载状态标签
        self.loading_label = QLabel("正在加载模型>>>>>>>>>")
        self.loading_label.setAlignment(Qt.AlignCenter)
        self.loading_label.setStyleSheet("font-size: 14px; color: blue;")
        self.loading_label.setVisible(True)
        main_layout.addWidget(self.loading_label)

        # 用于控制闪烁动画的变量
        self.loading_animation_active = True
        self.loading_dots_count = 8
        self.start_loading_animation()


        # 第二行：Excel导出选项
        export_group_box = QWidget()
        export_layout = QVBoxLayout(export_group_box)
        export_label = QLabel("excel导出选项(直接包含图片文件最直观但excel体积会变的很大)")
        export_layout.addWidget(export_label)

        # 单选按钮组
        radio_layout = QHBoxLayout()
        self.radio_group = QButtonGroup()

        self.radio_image_path = QRadioButton("包含图片路径")
        self.radio_image_file = QRadioButton("包含图片文件")
        self.radio_image_file.setChecked(True)  # 默认选中

        self.radio_group.addButton(self.radio_image_path)
        self.radio_group.addButton(self.radio_image_file)

        radio_layout.addWidget(self.radio_image_path)
        radio_layout.addWidget(self.radio_image_file)
        export_layout.addLayout(radio_layout)

        main_layout.addWidget(export_group_box)

        # 第三行：重命名选项
        rename_group_box = QWidget()
        rename_layout = QVBoxLayout(rename_group_box)
        rename_label = QLabel("重命名选项")
        rename_layout.addWidget(rename_label)

        # 主控制单选框
        rename_control_layout = QHBoxLayout()
        self.radio_no_rename = QRadioButton("不重命名")
        self.radio_do_rename = QRadioButton("重命名文件")
        self.radio_do_rename.setChecked(True)  # 默认启用重命名
        rename_control_layout.addWidget(self.radio_no_rename)
        rename_control_layout.addWidget(self.radio_do_rename)
        rename_layout.addLayout(rename_control_layout)

        # 重命名具体选项（直接添加到rename_layout，而不是通过中间widget）
        # 复选框选项
        checkbox_layout = QHBoxLayout()
        self.check_name = QCheckBox("姓名")
        self.check_sex = QCheckBox("性别")
        self.check_id = QCheckBox("身份证号")
        self.check_nation = QCheckBox("民族")
        self.check_address = QCheckBox("住址(省)")

        # 默认选中所有选项
        self.check_name.setChecked(True)
        self.check_id.setChecked(True)

        checkbox_layout.addWidget(self.check_name)
        checkbox_layout.addWidget(self.check_id)
        checkbox_layout.addWidget(self.check_nation)
        checkbox_layout.addWidget(self.check_sex)
        checkbox_layout.addWidget(self.check_address)
        checkbox_layout.addStretch()

        rename_layout.addLayout(checkbox_layout)

        # 分隔符输入框
        separator_layout = QHBoxLayout()
        separator_label = QLabel("分隔符:")
        self.separator_input = QLineEdit("-")  # 默认为_
        self.separator_input.setPlaceholderText("分隔符最好为 _ - + | 中的任意一个")
        separator_layout.addWidget(separator_label)
        separator_layout.addWidget(self.separator_input)
        separator_layout.addStretch()

        rename_layout.addLayout(separator_layout)

        # 连接信号槽，控制重命名选项的启用状态
        self.radio_no_rename.toggled.connect(self.toggle_rename_options)
        self.toggle_rename_options()  # 初始化状态

        main_layout.addWidget(rename_group_box)
        bottom_layout = QHBoxLayout()

        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 0)  # 设置为不确定模式
        self.progress_bar.hide()

        self.start_button = QPushButton("开始识别")
        self.start_button.clicked.connect(self.start_recognition)
        # self.start_button.setEnabled(False)  # 可以考虑初始状态下禁用

        # 添加终止按钮
        self.terminate_button = QPushButton("终止处理")
        self.terminate_button.clicked.connect(self.terminate_processing)
        self.terminate_button.setEnabled(False)  # 初始状态下禁用


        bottom_layout.addWidget(self.progress_bar)
        bottom_layout.addWidget(self.start_button)
        bottom_layout.addWidget(self.terminate_button)
        main_layout.addLayout(bottom_layout)

        main_layout.addStretch()

    def on_model_loading_finished(self, success, message):
        """模型加载完成后的回调"""
        # 停止动画并更新文本
        self.stop_loading_animation()

        if success:
            self.loading_label.setText("模型加载完毕！可以开始识别了！")
            self.loading_label.setStyleSheet("font-size: 14px; color: green; font-weight: bold;")
            print(message)
            QMessageBox.information(self, "提示", "模型加载完成，可以开始识别了！")
            # self.start_button.setEnabled(True)  # 启用开始识别按钮
        else:
            self.loading_label.setText("模型加载失败！")
            self.loading_label.setStyleSheet("font-size: 14px; color: red; font-weight: bold;")
            QMessageBox.warning(self, "警告", message)

    def start_recognition(self):
        """开始识别按钮的点击事件"""
        # 检查模型是否已加载完成
        if not self.model_loader.model_loaded:
            QMessageBox.information(self, "提示", "请稍等，正在加载模型...")
            return

            # 使用已加载的模型
        ocr_model = self.model_loader.ocr

        # 弹出文件选择框
        file_paths, _ = QFileDialog.getOpenFileNames(
            self,
            "选择身份证图片",
            "",
            "Image Files (*.jpg *.jpeg *.png *.bmp)"
        )

        if not file_paths:
            return

        # 显示进度条并开始处理
        self.progress_bar.show()
        self.start_button.setEnabled(False)
        self.terminate_button.setEnabled(True)  # 启用终止按钮

        # 创建并启动OCR工作线程
        export_options = self.get_export_options()
        self.ocr_worker = OCRWorker(file_paths, export_options, ocr_model)
        self.ocr_worker.progress_updated.connect(self.update_progress)
        self.ocr_worker.finished_signal.connect(self.on_processing_finished)
        self.ocr_worker.error_occurred.connect(self.on_processing_error)
        self.ocr_worker.start()


    def update_progress(self, current, total):
        """更新进度条"""
        self.progress_bar.setRange(0, total)
        self.progress_bar.setValue(current)

    def on_processing_error(self, error_message):
        """处理错误"""
        self.progress_bar.hide()
        self.start_button.setEnabled(True)
        QMessageBox.critical(self, "错误", error_message)

    def get_export_options(self):
        """获取导出选项"""
        export_option = "no_image"
        if self.radio_image_path.isChecked():
            export_option = "image_path"
        elif self.radio_image_file.isChecked():
            export_option = "image_file"

        # 检查是否启用重命名
        rename_enabled = self.radio_do_rename.isChecked()
        rename_options = []

        if rename_enabled:
            if self.check_name.isChecked():
                rename_options.append("name")
            if self.check_id.isChecked():
                rename_options.append("id")
            if self.check_nation.isChecked():
                rename_options.append("nation")
            if self.check_sex.isChecked():
                rename_options.append("sex")
            if self.check_address.isChecked():
                rename_options.append("address")

        separator = self.separator_input.text()

        return {
            "export_option": export_option,
            "rename_enabled": rename_enabled,
            "rename_options": rename_options,
            "separator": separator
        }

    def on_processing_finished(self):
        """处理完成后的回调"""
        self.progress_bar.hide()
        self.start_button.setEnabled(True)
        QMessageBox.information(self, "完成", "识别完成，结果已保存至 身份证识别结果.xlsx")

    def start_loading_animation(self):
        """开始加载动画 - 修复版循环滚动效果"""
        if self.loading_animation_active:
            # 实现循环滚动效果
            current_text = self.loading_label.text()
            if "正在加载模型" in current_text:
                # 提取箭头部分（去掉"正在加载模型"前缀）
                arrows_part = current_text.replace("正在加载模型", "", 1)

                # 确保有足够的箭头字符
                if not arrows_part:
                    arrows_part = ">>>>>>>>"

                # 循环移动箭头：将第一个字符移到末尾
                new_arrows = arrows_part[1:] + arrows_part[0] if len(arrows_part) > 1 else arrows_part

                # 更新标签文本
                self.loading_label.setText(f"正在加载模型{new_arrows}")

                # 颜色闪烁效果
                current_color = self.loading_label.styleSheet()
                if "color: blue;" in current_color:
                    self.loading_label.setStyleSheet("font-size: 14px; color: black; font-weight: normal;")
                else:
                    self.loading_label.setStyleSheet("font-size: 14px; color: blue; font-weight: normal;")

            # 每200毫秒更新一次动画
            self.loading_timer = self.startTimer(200)

    def timerEvent(self, event):
        """定时器事件，用于更新加载动画"""
        if hasattr(self, 'loading_timer') and event.timerId() == self.loading_timer:
            self.start_loading_animation()

    def stop_loading_animation(self):
        """停止加载动画"""
        self.loading_animation_active = False
        if hasattr(self, 'loading_timer'):
            self.killTimer(self.loading_timer)


    # 终止识别线程
    def terminate_processing(self):
        """终止处理按钮的点击事件"""
        if hasattr(self, 'ocr_worker') and self.ocr_worker.isRunning():
            # 弹出确认对话框
            reply = QMessageBox.question(
                self,
                "确认终止",
                "确定要终止处理吗？已处理的数据将会保存。",
                QMessageBox.Yes | QMessageBox.No,
                QMessageBox.No
            )

            if reply == QMessageBox.Yes:
                # 请求终止处理
                self.ocr_worker.request_termination()
                self.terminate_button.setEnabled(False)
                self.start_button.setEnabled(False)
                # 更新界面提示
                self.progress_bar.setFormat("正在保存已处理数据...")

    def toggle_rename_options(self):
        """根据主单选框状态控制重命名选项的启用/禁用"""
        enabled = self.radio_do_rename.isChecked()
        # 直接控制各个控件的启用状态
        self.check_name.setEnabled(enabled)
        self.check_sex.setEnabled(enabled)
        self.check_id.setEnabled(enabled)
        self.check_nation.setEnabled(enabled)
        self.check_address.setEnabled(enabled)
        self.separator_input.setEnabled(enabled)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = StartupWindow()
    window.show()
    sys.exit(app.exec())
